//
//! Copyright 2020 Alibaba Group Holding Limited.
//! 
//! Licensed under the Apache License, Version 2.0 (the "License");
//! you may not use this file except in compliance with the License.
//! You may obtain a copy of the License at
//! 
//!     http://www.apache.org/licenses/LICENSE-2.0
//! 
//! Unless required by applicable law or agreed to in writing, software
//! distributed under the License is distributed on an "AS IS" BASIS,
//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//! See the License for the specific language governing permissions and
//! limitations under the License.

use super::*;
use crate::operator::unary::Unary;
use crate::communication::Pipeline;

pub trait Inspect<D: Data, A> {
    fn inspect<F>(&self, func: F) -> Stream<D, A> where F: Fn(&D) + Send + 'static;

    fn inspect_batch<F>(&self, func: F) -> Stream<D, A> where F: Fn(&[D]) + Send + 'static;
}

impl<D: Data, A: DataflowBuilder> Inspect<D, A> for Stream<D, A> {
    fn inspect<F>(&self, func: F) -> Stream<D, A> where F: Fn(&D) + Send + 'static {
        self.unary("inspect", Pipeline, |info| {
            info.set_pass();
            move |input, output| {
                input.for_each_batch(|dataset| {
                    let mut session = output.session(&dataset);
                    let data = dataset.data();
                    for datum in data.iter() {
                        func(datum)
                    }
                    session.give_batch(data)?;
                    Ok(true)
                })?;
                Ok(())
            }
        })
    }

    fn inspect_batch<F>(&self, func: F) -> Stream<D, A> where F: Fn(&[D]) + Send + 'static {
        self.unary("inspect_batch", Pipeline, |info| {
            info.set_pass();
            move |input, output| {
                input.for_each_batch(|dataset| {
                    let mut session = output.session(&dataset);
                    let data = dataset.data();
                    func(&data);
                    session.give_batch(data)?;
                    Ok(true)
                })?;
                Ok(())
            }
        })
    }
}
